
import cv2
import time
import random
import numpy as np


def adjust_brightness_contrast_saturation(img, brightness=0, contrast=1.0, saturation=1.0):
    # 调整亮度与对比度
    adjusted = cv2.convertScaleAbs(img, alpha=contrast, beta=brightness)
    # 调整饱和度
    hsv = cv2.cvtColor(adjusted, cv2.COLOR_BGR2HSV).astype(np.float32)
    hsv[...,1] = np.clip(hsv[...,1] * saturation, 0, 255)
    adjusted_hsv = cv2.cvtColor(hsv.astype(np.uint8), cv2.COLOR_HSV2BGR)
    return adjusted_hsv

def main():
    cap = cv2.VideoCapture("/dev/video_4060h")
    if not cap.isOpened():
        print("无法打开摄像头")
        return

    cv2.namedWindow("Capture", cv2.WINDOW_NORMAL)
    paused = False
    frame_index = 1
    last_time = time.time()

    while True:
        ret, frame = cap.read()
        if not ret:
            break

        cv2.imshow("Capture", frame)
        key = cv2.waitKey(1) & 0xFF

        # 按 z/Z 暂停或继续
        if key in [ord('z'), ord('Z')]:
            paused = not paused
            if paused:
                print("NOT SHOTED")
            else:
                print("RECORDING")
        
        
        # 按 ESC/Q 退出
        if key == 27 or key == ord('q'):
            break

        # 每0.5秒采集一次并保存
        if not paused and (time.time() - last_time) >= 0.5:
            last_time = time.time()

            # 生成13张变换图 （包含原图在内的各种组合）
            for i in range(1, 14):
                # 随机设置亮度、对比度、饱和度参数
                b_offset = random.randint(-50, 50)
                c_factor = round(random.uniform(0.5, 1.5), 2)
                s_factor = round(random.uniform(0.5, 1.5), 2)

                transformed = adjust_brightness_contrast_saturation(frame, b_offset, c_factor, s_factor)
                # 文件名形如：img0001xx.png，其中0001 表示第1次采集，xx表示第几种变换
                filename = f"/home/firefly/Documents/yolo_recorded/images/img{frame_index:04d}{i:02d}.png"
                cv2.imwrite(filename, transformed)

            frame_index += 1

    cap.release()
    cv2.destroyAllWindows()

if __name__ == "__main__":
    main()
